package com.um.springboot.starter.interceptors;

import com.um.springboot.starter.annotation.RateLimiter;
import com.um.springboot.starter.orm.entity.Response;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.annotation.Order;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpSession;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 流量限量拦截器
 * @author 邓南辉
 * @date 2019/12/13
 * @copyright 广州优匠科技有限公司
 */
@Slf4j
@Aspect
@Component
@Order(value = AspectOrders.LIMITER_ORDER)
public class RateLimiterInterceptor {

    @Autowired
    private HttpServletRequest request;

    @Autowired
    private StringRedisTemplate redis;

    @Around("@annotation(com.um.springboot.starter.annotation.RateLimiter)")
    public Object doAround(ProceedingJoinPoint pjp) throws Throwable {

        MethodSignature signature = (MethodSignature) pjp.getSignature();
        RateLimiter rateLimiter = signature.getMethod().getAnnotation(RateLimiter.class);

        //有注解的才需要做限流验证
        if(rateLimiter != null) {
            String key = rateLimiter.type() == RateLimiter.LimiterType.IP ? getKeyByIp() : getKeyBySession();

            int timeout = rateLimiter.timeout();

            long times = incrementAndGet(rateLimiter.store(),key,timeout);
            if (times > rateLimiter.max()) {
                String message = "请求过于频繁,请" + timeout + " " + "分钟后再试!";
                return Response.failed(message);
            }

        }

        //result的值就是被拦截方法的返回值
        return pjp.proceed();
    }

    private String getKeyBySession() {
        return "api:" + request.getSession().getId()+ ":" + request.getRequestURI();
    }

    private String getKeyByIp() {
        String ip = request.getRemoteAddr();
        return "api:" + ip + ":" + request.getRequestURI();
    }

    /**
     * 增加1，并返回当前的访问次量
     * @param storeType 存储类型
     * @param key 键
     * @param timeout 过期时间
     * @return
     */
    private long incrementAndGet(RateLimiter.StoreType storeType,String key,int timeout){
        if(storeType == RateLimiter.StoreType.REDIS) {
            boolean hasKey = redis.hasKey(key);
            Long value = redis.opsForValue().increment(key, 1);
            if (!hasKey) {
                redis.expire(key, timeout, TimeUnit.MINUTES);
            }
            return value;
        }else{

            HttpSession session = request.getSession();
            List<Date> rateList = (List<Date>)session.getAttribute(key);
            if(rateList == null){
                rateList = new ArrayList<>();
            }

            Date currentDate = new Date();

            //删除过期访问
            int second = timeout * 60;
            for(int index = rateList.size() -1;index >= 0;index --){
                Date date = rateList.get(index);
                long diff = (currentDate.getTime() - date.getTime())/1000;
                if(diff > second){
                    rateList.remove(index);
                }
            }

            //添加本次访问时间
            rateList.add(currentDate);

            session.setAttribute(key,rateList);
            return rateList.size();

        }
    }
}
